import torch
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from datasets import load_dataset

# 1. 加载数据集
dataset = load_dataset("imdb")
train_dataset = dataset["train"]
test_dataset = dataset

# 2. 加载BERT tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# 3. 数据预处理
def preprocess_function(examples):
    return tokenizer(examples['text'], truncation=True, padding=True)


trainodings = preprocess_function(train_dataset)
test_encodings = preprocess_function(test_dataset)


# 4. 创建PyTorch数据集
class IMDbDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)


train_labels = train_dataset['label']
test_labels = test_dataset['label']


